import numpy as np
import scipy.spatial.transform as st


# ====== start: tf utils ======
def pos_rot_to_mat(pos, rot):
    shape = pos.shape[:-1]
    mat = np.zeros(shape + (4,4), dtype=pos.dtype)
    mat[...,:3,3] = pos
    mat[...,:3,:3] = rot.as_matrix()
    mat[...,3,3] = 1
    return mat

def mat_to_pos_rot(mat):
    pos = (mat[...,:3,3].T / mat[...,3,3].T).T
    rot = st.Rotation.from_matrix(mat[...,:3,:3])
    return pos, rot

def pos_rot_to_pose(pos, rot):
    shape = pos.shape[:-1]
    pose = np.zeros(shape+(6,), dtype=pos.dtype)
    pose[...,:3] = pos
    pose[...,3:] = rot.as_rotvec()
    return pose

def pose_to_pos_rot(pose):
    pos = pose[...,:3]
    rot = st.Rotation.from_rotvec(pose[...,3:])
    return pos, rot

def pose_to_mat(pose):
    return pos_rot_to_mat(*pose_to_pos_rot(pose))

def mat_to_pose(mat):
    return pos_rot_to_pose(*mat_to_pos_rot(mat))

def transform_pose(tx, pose):
    """
    tx: tx_new_old
    pose: tx_old_obj
    result: tx_new_obj
    """
    pose_mat = pose_to_mat(pose)
    tf_pose_mat = tx @ pose_mat
    tf_pose = mat_to_pose(tf_pose_mat)
    return tf_pose

def transform_point(tx, point):
    return point @ tx[:3,:3].T + tx[:3,3]

def project_point(k, point):
    x = point @ k.T
    uv = x[...,:2] / x[...,[2]]
    return uv

def apply_delta_pose(pose, delta_pose):
    new_pose = np.zeros_like(pose)

    # simple add for position
    new_pose[:3] = pose[:3] + delta_pose[:3]

    # matrix multiplication for rotation
    rot = st.Rotation.from_rotvec(pose[3:])
    drot = st.Rotation.from_rotvec(delta_pose[3:])
    new_pose[3:] = (drot * rot).as_rotvec()

    return new_pose

def normalize(vec, tol=1e-7):
    return vec / np.maximum(np.linalg.norm(vec), tol)

def rot_from_directions(from_vec, to_vec):
    from_vec = normalize(from_vec)
    to_vec = normalize(to_vec)
    axis = np.cross(from_vec, to_vec)
    axis = normalize(axis)
    angle = np.arccos(np.dot(from_vec, to_vec))
    rotvec = axis * angle
    rot = st.Rotation.from_rotvec(rotvec)
    return rot

def normalize(vec, eps=1e-12):
    norm = np.linalg.norm(vec, axis=-1)
    norm = np.maximum(norm, eps)
    out = (vec.T / norm).T
    return out

def rot6d_to_mat(d6):
    a1, a2 = d6[..., :3], d6[..., 3:]
    b1 = normalize(a1)
    b2 = a2 - np.sum(b1 * a2, axis=-1, keepdims=True) * b1
    b2 = normalize(b2)
    b3 = np.cross(b1, b2, axis=-1)
    out = np.stack((b1, b2, b3), axis=-2)
    return out

def mat_to_rot6d(mat):
    batch_dim = mat.shape[:-2]
    out = mat[..., :2, :].copy().reshape(batch_dim + (6,))
    return out

def mat_to_pose10d(mat):
    pos = mat[...,:3,3]
    rotmat = mat[...,:3,:3]
    d6 = mat_to_rot6d(rotmat)
    d10 = np.concatenate([pos, d6], axis=-1)
    return d10

def pose10d_to_mat(d10):
    pos = d10[...,:3]
    d6 = d10[...,3:]
    rotmat = rot6d_to_mat(d6)
    out = np.zeros(d10.shape[:-1]+(4,4), dtype=d10.dtype)
    out[...,:3,:3] = rotmat
    out[...,:3,3] = pos
    out[...,3,3] = 1
    return out

def compute_relative_pose(pos, rot, base_pos, base_rot_mat,
                            rot_transformer_to_mat,
                            rot_transformer_to_target,
                            backward=False,
                            delta=False):
    if not backward:
        # forward pass
        if not delta:
            output_pos = pos if base_pos is None else pos - base_pos
            output_rot = rot_transformer_to_target.forward(
                rot_transformer_to_mat.forward(rot) @ np.linalg.inv(base_rot_mat))
            return output_pos, output_rot
        else:
            all_pos = np.concatenate([base_pos[None,...], pos], axis=0)
            output_pos = np.diff(all_pos, axis=0)
            
            rot_mat = rot_transformer_to_mat.forward(rot)
            all_rot_mat = np.concatenate([base_rot_mat[None,...], rot_mat], axis=0)
            prev_rot = np.linalg.inv(all_rot_mat[:-1])
            curr_rot = all_rot_mat[1:]
            rot = np.matmul(curr_rot, prev_rot)
            output_rot = rot_transformer_to_target.forward(rot)
            return output_pos, output_rot
            
    else:
        # backward pass
        if not delta:
            output_pos = pos if base_pos is None else pos + base_pos
            output_rot = rot_transformer_to_mat.inverse(
                rot_transformer_to_target.inverse(rot) @ base_rot_mat)
            return output_pos, output_rot
        else:
            output_pos = np.cumsum(pos, axis=0) + base_pos
            
            rot_mat = rot_transformer_to_target.inverse(rot)
            output_rot_mat = np.zeros_like(rot_mat)
            curr_rot = base_rot_mat
            for i in range(len(rot_mat)):
                curr_rot = rot_mat[i] @ curr_rot
                output_rot_mat[i] = curr_rot
            output_rot = rot_transformer_to_mat.inverse(rot)
            return output_pos, output_rot


def convert_pose_mat_rep(pose_mat, base_pose_mat, pose_rep='abs', backward=False):
    if not backward:
        # training transform
        if pose_rep == 'abs':
            return pose_mat
        elif pose_rep == 'rel':
            # legacy buggy implementation
            # for compatibility
            pos = pose_mat[...,:3,3] - base_pose_mat[:3,3]
            rot = pose_mat[...,:3,:3] @ np.linalg.inv(base_pose_mat[:3,:3])
            out = np.copy(pose_mat)
            out[...,:3,:3] = rot
            out[...,:3,3] = pos
            return out
        elif pose_rep == 'relative':
            out = np.linalg.inv(base_pose_mat) @ pose_mat
            return out
        elif pose_rep == 'delta':
            all_pos = np.concatenate([base_pose_mat[None,:3,3], pose_mat[...,:3,3]], axis=0)
            out_pos = np.diff(all_pos, axis=0)
            
            all_rot_mat = np.concatenate([base_pose_mat[None,:3,:3], pose_mat[...,:3,:3]], axis=0)
            prev_rot = np.linalg.inv(all_rot_mat[:-1])
            curr_rot = all_rot_mat[1:]
            out_rot = np.matmul(curr_rot, prev_rot)
            
            out = np.copy(pose_mat)
            out[...,:3,:3] = out_rot
            out[...,:3,3] = out_pos
            return out
        else:
            raise RuntimeError(f"Unsupported pose_rep: {pose_rep}")

    else:
        # eval transform
        if pose_rep == 'abs':
            return pose_mat
        elif pose_rep == 'rel':
            # legacy buggy implementation
            # for compatibility
            pos = pose_mat[...,:3,3] + base_pose_mat[:3,3]
            rot = pose_mat[...,:3,:3] @ base_pose_mat[:3,:3]
            out = np.copy(pose_mat)
            out[...,:3,:3] = rot
            out[...,:3,3] = pos
            return out
        elif pose_rep == 'relative':
            out = base_pose_mat @ pose_mat
            return out
        elif pose_rep == 'delta':
            output_pos = np.cumsum(pose_mat[...,:3,3], axis=0) + base_pose_mat[:3,3]
            
            output_rot_mat = np.zeros_like(pose_mat[...,:3,:3])
            curr_rot = base_pose_mat[:3,:3]
            for i in range(len(pose_mat)):
                curr_rot = pose_mat[i,:3,:3] @ curr_rot
                output_rot_mat[i] = curr_rot
            
            out = np.copy(pose_mat)
            out[...,:3,:3] = output_rot_mat
            out[...,:3,3] = output_pos
            return out
        else:
            raise RuntimeError(f"Unsupported pose_rep: {pose_rep}")
        
def convert_iterative_action10d_to_0relative(action_pose10d):
    """
    将依次相对前一时刻的动作转换为相对于第0维的动作。
    pose使用3维pos+6维rot的表示方式。
    Args:
        actions (np.ndarray): 形状为 (N, 9) 的数组。

    Returns:
        np.ndarray: 形状为 (N, 9) 的数组，表示相对于第0维的动作。
    """
    action_pose10d_0relative = np.zeros_like(action_pose10d)
    current_pose = np.eye(4)
    for i in range(action_pose10d.shape[0]):
        
        action_pose10d_i = action_pose10d[i]
        action_pose_mat_i = pose10d_to_mat(action_pose10d_i)

        action_mat_i = convert_pose_mat_rep(action_pose_mat_i, 
                                            base_pose_mat=current_pose,
                                            pose_rep='relative',
                                            backward=True)
        
        action_pose10d_0relative[i] = mat_to_pose10d(action_mat_i)
        current_pose = action_mat_i.copy()
    
    return action_pose10d_0relative
# ====== end: tf utils ======

# ====== start: sample utils ======
def sample_chunk(data, idx_current, idx_relate, idx_padding=-1):
    """
    从 data 数组中沿第 0 维采样一个数据块。

    Args:
        data (np.ndarray): 要从中采样的 N 维数组。采样沿第 0 维进行。
        idx_current (int): 当前的中心索引（沿第 0 维）。
        idx_relate (list or np.ndarray): 相对于 idx_current 的索引列表。
        idx_padding (int): 当采样索引超出 data 第 0 维边界时，
                            用于填充的 idx_relate 中的索引。
                            默认为 -1，表示使用 idx_relate 中的最后一个元素。

    Returns:
        np.ndarray: 采样得到的数据块。形状为 (len(idx_relate),) + data.shape[1:]。
    """
    idx_relate = np.array(idx_relate)
    absolute_indices = idx_current + idx_relate
    # 获取第 0 维的长度
    data_len = data.shape[0]

    # 计算用于填充的绝对索引（沿第 0 维）
    padding_absolute_idx = idx_current + idx_relate[idx_padding]
    # 将填充索引限制在 data 第 0 维的有效范围内
    padding_absolute_idx_clamped = np.clip(padding_absolute_idx, 0, data_len - 1)
    # 获取填充值（沿第 0 维索引，保留其他维度）
    padding_value_slice = data[padding_absolute_idx_clamped, ...]

    # 计算结果数组的形状
    result_shape = (len(absolute_indices),) + data.shape[1:]
    # 创建结果数组，并用填充值切片初始化
    # 注意：这里不能直接用 fill_value=padding_value_slice，因为 fill_value 期望标量
    # 我们先创建一个空数组，然后用广播填充
    sampled_chunk = np.empty(result_shape, dtype=data.dtype)
    sampled_chunk[:] = padding_value_slice # 使用广播填充

    # 找出有效的索引（在 data 第 0 维边界内）
    valid_mask = (absolute_indices >= 0) & (absolute_indices < data_len)

    # 获取有效的绝对索引（沿第 0 维）
    valid_absolute_indices = absolute_indices[valid_mask]

    # 从 data 中获取有效索引对应的值（沿第 0 维索引，保留其他维度）
    valid_values = data[valid_absolute_indices, ...]

    # 将有效值放入结果数组的相应位置（沿第 0 维赋值）
    sampled_chunk[valid_mask, ...] = valid_values

    return sampled_chunk

def sample_history(array, idx, chunk_size):
    # 计算起始索引
    start_idx = max(0, idx - chunk_size + 1)
    
    # 从起始索引到idx的部分
    chunk = array[start_idx:idx+1]
    
    # 如果chunk的长度不足chunk_size，填充array[0]的值
    padding_size = chunk_size - len(chunk)
    if padding_size > 0:
        padding = np.tile(array[0], (padding_size, 1))  # 用array[0]填充
        chunk = np.vstack([padding, chunk])  # 将填充和现有的部分合并
    
    # 将chunk展平成1维
    return chunk.flatten()
# ====== end: sample utils ======